/*
 *  Copyright 2009 Advanced Industrial Science and Technology (AIST).
 *
 *  Licensed under the Apache License, Version 2.0 (the "License");
 *  you may not use this file except in compliance with the License.
 *  You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 *  Unless required by applicable law or agreed to in writing, software
 *  distributed under the License is distributed on an "AS IS" BASIS,
 *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *  See the License for the specific language governing permissions and
 *  limitations under the License.
 */

/*! \file
 *  \brief
 */

#ifdef __CUDACC__

// stl, includes
#include <algorithm>

// thrust, includes
#include <thrust/iterator/iterator_traits.h>
#include <thrust/detail/raw_buffer.h>
#include <thrust/fill.h>
#include <thrust/count.h>

#include <thrust/detail/device/cuda/arch.h>
#include <thrust/detail/device/dereference.h>

#include <thrust/graph/properties.hpp>

#include <iostream>

namespace thrust {

namespace detail {

namespace device {
      
namespace cuda {

template <typename Graph
          , typename VertexIterator
          , typename BFSVisitor>
__global__
void bfs_initialize_kernel( Graph g,
                            VertexIterator first,
                            VertexIterator last, 
                            BFSVisitor vis)
{
  typedef typename graph_traits<Graph>::vertex_descriptor vertex_descriptor;
  typedef typename thrust::iterator_traits<VertexIterator>::difference_type IndexType;
  
  const IndexType grid_size =blockDim.x * gridDim.x;
  
  VertexIterator iter =first + blockIdx.x * blockDim.x + threadIdx.x;
  
  while( iter <last){
    vertex_descriptor u =*iter;
    vis.initialize_vertex( u);
    iter +=grid_size;
  }
}

template <typename Graph
          , typename VertexIterator
          , typename BFSVisitor
          >
__global__
void bfs_propagation_kernel( Graph g,
                             VertexIterator first,
                             VertexIterator last,
                             int* current_color,
                             int* buffer_color,
                             BFSVisitor vis
                             )
{
  typedef typename graph_traits<Graph>::vertex_descriptor   vertex_descriptor;
  typedef typename graph_traits<Graph>::vertices_size_type  vertices_size_type;
  typedef typename graph_traits<Graph>::degree_size_type    degree_size_type;
  typedef typename graph_traits<Graph>::kernel_adjacency_iterator kernel_adjacency_iterator;

  typedef typename thrust::iterator_traits<VertexIterator>::difference_type IndexType;

  const IndexType grid_size =blockDim.x * gridDim.x;
  
  VertexIterator iter =first + blockIdx.x * blockDim.x + threadIdx.x;

  iter =first + blockIdx.x * blockDim.x + threadIdx.x;
  
  while( iter <last){

    vertex_descriptor cur_vertex =*iter;
    vis.enter_step( cur_vertex);
    if( current_color[cur_vertex] ==gray_color){
      
      kernel_adjacency_iterator adj_iter, adj_last;
      tie( adj_iter, adj_last)
        =adjacent_vertices( cur_vertex, g, thrust::called_from_device_tag());
      
      for( ; adj_iter !=adj_last; ++adj_iter){
        vertex_descriptor next_vertex =*adj_iter;
        if( current_color[next_vertex] ==white_color){
          buffer_color[next_vertex] =gray_color;
          vis.discover_vertex( next_vertex); // discover adjacent vertex v
        }
        else{
          
        }
      }
      buffer_color[cur_vertex] =black_color;
      vis.finish_vertex( cur_vertex); // finish
    }
    vis.out_step( cur_vertex); //out_step
    
    iter +=grid_size;
  } // end while( iter <last)
}

struct color_equal_gray
{
  __device__
  inline bool
  operator()( const int& x) const
  {
    return x ==gray_color;
  }
};

template <class Graph, class Visitor>
void 
breadth_first_search( const Graph& g
                      , typename graph_traits<Graph>::vertex_descriptor s
                      , Visitor& vis)
{
  typedef typename graph_traits<Graph>::vertex_iterator vertex_iterator;
  vertex_iterator first, last;
  tie( first, last) =vertices( g);
  
  // initilize color white_color
  typedef color_traits<int> Color;
  thrust::detail::raw_cuda_device_buffer<int> current_color( num_vertices(g));
  thrust::detail::raw_cuda_device_buffer<int> buffer_color( num_vertices(g));
  
  thrust::fill( current_color.begin(), current_color.end(), white_color);
  thrust::fill( buffer_color.begin(), buffer_color.end(), white_color);
  
  current_color[s] =gray_color;
  
  int* ptr_current_color =thrust::raw_pointer_cast(&current_color[0]);
  int* ptr_buffer_color  =thrust::raw_pointer_cast(&buffer_color[0]);
  
  if( first >=last) return;
  const std::size_t BLOCK_SIZE =256; // number of threads per block
  const std::size_t MAX_BLOCKS =thrust::detail::device::cuda::arch::max_active_threads() / BLOCK_SIZE;
  const std::size_t NUM_BLOCKS =std::min( MAX_BLOCKS, ((last - first) + (BLOCK_SIZE - 1)) / BLOCK_SIZE);
  
#ifdef BFS_DEBUG
  std::cout <<"(MAX_BLOCKS:" <<MAX_BLOCKS
            <<", NUM_BLOCKS:" <<NUM_BLOCKS
            <<", BLOCK_SIZE:" <<BLOCK_SIZE
            <<") " <<std::endl;
  
  for( std::size_t i =0; i <num_vertices(g); ++i){
    std::cout <<"color["
              <<i <<"]:"
              <<current_color[i] 
              <<std::endl;
  }
  std::cout <<std::endl;
#endif
	      
  bfs_initialize_kernel<<<NUM_BLOCKS, BLOCK_SIZE>>>( g, first, last, vis);
  
  int loop =0;
  int count =thrust::count_if( current_color.begin(), 
                               current_color.end(),
                               color_equal_gray());
  while( count){
    ++loop;
    bfs_propagation_kernel<<<NUM_BLOCKS, BLOCK_SIZE>>>( g,             // Graph
                                                        first,         // VertexIterator
                                                        last,          // VertexIterator
                                                        ptr_current_color, // Current Color Map
                                                        //current_color.begin().base(),
                                                        ptr_buffer_color,  // Buffere Color Map
                                                        //buffer_color.begin().base(),
                                                        vis            // BFS Visitor
                                                        );
    
    thrust::copy( buffer_color.begin(), 
                  buffer_color.end(), 
                  current_color.begin());
    
    count =thrust::count_if( current_color.begin(), 
                             current_color.end(),
                             color_equal_gray());
    //if( loop ==1)
    //break;
  } // end while( count)*/

#ifdef BFS_DEBUG
  for( std::size_t i =0; i <num_vertices(g); ++i)
    std::cout <<"color[" 
              <<i <<"]:"
              <<current_color[i] 
              <<std::endl;
#endif
  //std::cout <<"loop: " <<loop <<std::endl;
}
  
} // end namespace cuda
      
} // end namespace device
  
} // end namespace detail
  
} // end namespace thrust
#endif // __CUDACC__
